#!/usr/bin/env python
###############################################################################
# VMWrapper.py: Python wrapper for BOINC-VM Apps
###############################################################################
# Copyright (C) 2009-10 CERN
#
# Authors: Jarno Rantala <jarno.rantala AT gmail.com>
#          Rohit Yadav <rohityadav89 AT gmail.com>

#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
###############################################################################

import os, signal, string, subprocess, sys, time, traceback, wx, xmlrpclib
from xml.dom import minidom
from boinc   import *
from gui     import *
from util    import *

sys.path.append(os.path.abspath('.'))

JOB_FILENAME = "job.xml"
CHECKPOINT_FILENAME  = "checkpoint.txt"
SERVER_PROXY = 'http://localhost:8090'

POLL_PERIOD = 1.0
CHECKPOINT_PERIOD = 15*60
TRICKLE_UP_PERIOD = 120

MAX_VMC_START_TIME  = 10   # Max time it takes VMController to start
MAX_VM_BOOT_TIME    = 90   # Max time it takes the Virtual Machine to boot up
MAX_VM_RESTORE_TIME = 10   # Max time it takes the Virtual Machine to restore

TASK_TAGS = ['virtualmachine',
           'image_path',
           'application',
           'type',
           'copy_app_to_VM',
           'copy_file_to_VM',
           'stdin_filename',
           'stdout_filename',
           'stderr_filename',
           'copy_file_from_VM',
           'checkpoint_filename',
           'command_line',
           'weight']

class VMState:
    poweredOff = "PoweredOff"
    runnning   = "Running"
    paused     = "Paused"
    saved      = "Saved"

###############################################################################
# parseJobFile(): parses job.xml file and return task list
###############################################################################
def parseJobFile(filename):
   input_path = boinc_resolve_filename(filename)

   # open the job file  
   try:
       infile = boinc_fopen(input_path, 'r')
   except boinc.error:  
      sys.stderr.write("Can't open job file: " + input_path)   
      raise Exception(1)

   jobxml = minidom.parse(infile)
   infile.close()

   # read the context of job file
   try:  
      xmltasks = jobxml.getElementsByTagName("job_desc")[0]
   except IndexError:
      sys.stderr.write("Can't read job file: no 'job_desc' tag \n")
     
   tasks = []

   # read attributes of unzip tasks
   unzipTasks = []
   for xmltask in xmltasks.getElementsByTagName("unzip_task"):
      task = TASK()
     
      for tag in TASK_TAGS:
         try:
            taglist = xmltask.getElementsByTagName(tag)
            for tagxml in taglist:
               data = tagxml.childNodes[0].data
               task.readTag(tag, data)  
         except IndexError:
            sys.stderr.write("Task has no " + tag + " \n")

      task.taskType = TaskType.unzip_task
      task.debug()         
      unzipTasks.append(task)

   # read the attributes of tasks
   for xmltask in xmltasks.getElementsByTagName("task"):
      task = TASK()
     
      for tag in TASK_TAGS:
         try:
            taglist = xmltask.getElementsByTagName(tag)
            for tagxml in taglist:
               data = tagxml.childNodes[0].data
               task.readTag(tag, data)  
         except IndexError:
            sys.stderr.write("Task has no "+tag+" \n")

      task.taskType = TaskType.task
      task.debug()
      tasks.append(task)

   # read attributes of VMmanage tasks
   vmTasks = []
   for xmltask in xmltasks.getElementsByTagName("vm_task"):
      task = TASK()
     
      for tag in TASK_TAGS:
         try:
            taglist = xmltask.getElementsByTagName(tag)
            for tagxml in taglist:
               data = tagxml.childNodes[0].data
               task.readTag(tag, data)  
         except IndexError:
            sys.stderr.write("Task has no "+tag+" \n")    

      task.taskType = TaskType.vm_task
      task.debug()       
      vmTasks.append(task)

   jobxml.unlink()
   return [unzipTasks, tasks, vmTasks]

###############################################################################
# readCheckPoint(): Checks last checkpoint file
###############################################################################
def readCheckPoint(filename):
   ntasks = 0
   cpu = 0

   try:
      f = open(filename, "r");
   except:
      return [ntasks, cpu]

   data = f.readline().split()
   f.close()

   try:
      ntasks = int(data[0])
      cpu = float(data[1])
   except:
      sys.stderr.write("Can't read checkpoint file \n")
      return [0, 0]

   return [ntasks, cpu]   

###############################################################################
# writeCheckPoint(): Writes a checkpoint file
###############################################################################
# Support for multiple tasks.
# We keep a checkpoint file that says how many tasks we've completed
# and how much CPU time has been used so far
def writeCheckPoint(filename, ntasks, cpu):
   try:
      f = open(filename, "w")
   except IOError:
      sys.stderr.write("Writing checkpoint file failed. \n")

   f.write(str(ntasks)+" "+str(cpu)+"\n")
   f.close()

###############################################################################
# pollBoincMessages(): Polls boinc messages
###############################################################################
def pollBoincMessages(task, VMmanage):
   status = boinc_get_status()
   exit = 0 # if nonzero then VMwrapper should exit
   #sys.stderr.write("status suspended: "+str(status['suspended'])+" \n")
   #sys.stderr.write("status no heartbeat: "+str(status['no_heartbeat'])+" \n")
   #sys.stderr.write("status quit request: "+str(status['quit_request'])+" \n")
   #sys.stderr.write("status abort request: "+str(status['abort_request'])+" \n")

   if status['no_heartbeat'] or status['quit_request'] or status['abort_request']:
      task.kill(VMmanage)
      exit = 1

   if status['suspended']:
      if not task.suspended:
         task.stop(VMmanage)
   else:
      if task.suspended:
         task.resume(VMmanage)

   return exit

###############################################################################
# sendStatusMessage(): Sends status message to boinc...
###############################################################################
def send_status_message(task, VMmanage, frac_done, checkpoint_cpu_time):
   current_cpu_time = task.starting_cpu + task.cpu_time(VMmanage)
   boinc_report_app_status(current_cpu_time, checkpoint_cpu_time, frac_done)

#def wait(n):
#   tic = time.time()
#   toc = time.time()
#   while toc - tic < n:
#      time.sleep(1)
#      toc = time.time()  

###############################################################################
# Main(): Startpoint of the program
###############################################################################
def main():
    # BOINC OPTIONS (bools) :
    main_program = 1
    check_heartbeat = 1
    handle_process_control = 1
    send_status_msgs = 1
    handle_trickle_ups = 1
    handle_trickle_downs = 1

    retval = boinc_init_options(main_program, check_heartbeat, handle_process_control,
             send_status_msgs, handle_trickle_ups, handle_trickle_downs)

    if retval:
        sys.exit(retval)

    # Start GUI
    app = wx.App()
    ui = WrapperGUI(None, -1, 'BoincVM Wrapper')

    logger("VMWrapper", "Parsing job.xml file.")
    # Parse job.xml file, returns list of different tasks
    [unzip_tasks, tasks, vm_tasks] = parseJobFile(JOB_FILENAME)

    # TODO: In future add code to read CheckPoint file

    # Unzip the archives
    logger("Unzip Task", "Unzipping files")
    for task in unzip_tasks:
        task.run()
    logger("Unzip Task", "Finished unzip_tasks")

    # Execute normal tasks first
    for task in tasks:
        print '[Running task]: Host tasks'
        task.run()
    # Note: VMController is executed as a normal in-host task

    virtualMachine = TASK()
    # Get VM-Tasks
    for task in vm_tasks:
        if task.appType == "vm_run":
            virtualMachine = task

    logger("VM Name", virtualMachine.vmName)
    logger("Wrapper", "Waiting for VM Controller...")

    time.sleep(MAX_VMC_START_TIME)
    # VM Controller RPC object
    vmController = xmlrpclib.ServerProxy(SERVER_PROXY)

    try:
        vmList =  vmController.listAvailableVMs()
    except:
        logger("VMWrapper", "Could n't connect to VMController, exiting.")
        for task in tasks:   
            task.kill()
        sys.exit(1)

    logger("VM List", vmList)

    try:
        if not ( virtualMachine.vmName in vmList) :
            logger("VM", "Trying to create VM")
            vmController.createVM(virtualMachine.vmName, os.path.abspath(virtualMachine.vmImagePath))
    except:
        logger("VMController", "Failed to create the VM")
        for task in tasks:   
            task.kill()
        sys.exit(1) 

    vmState = "NULL"

    try:
        vmState = vmController.getState(virtualMachine.vmName)
        logger("VM", "Trying to start VM")
        vmController.start(virtualMachine.vmName)

        # If machine is poweredOff, we wait for the MAX_VM_BOOT_TIME
        if vmState == VMState.poweredOff:
            logger("WAITING", "Booting VM")
            time.sleep(MAX_VM_BOOT_TIME)
        else: 
            logger("WAITING", "Restoring VM")
            time.sleep(MAX_VM_RESTORE_TIME)
    except:
        logger("VMController", "Unable to start the VM")
        for task in tasks:   
            task.kill()
        sys.exit(1) 

    logger("VM", "Starting VM")

    while 1:
        try:
            print vmController.ping(virtualMachine.vmName)
            break
        except:
            logger("Wrapper", "Waiting for response from VMController in VM")
        time.sleep(1)

    # Infinite loop, stays until, wrapper ui is exited
    app.MainLoop()

    try:
        vmState = vmController.getState(virtualMachine.vmName)
        print "VMState", vmState

        if vmState == VMState.runnning:
            print "Trying to save the state of VM: ", vmController.saveState(virtualMachine.vmName)
        elif vmState == VMState.paused:
            vmController.unpause(virtualMachine.vmName)
            vmController.saveState(virtualMachine.vmName)
    except:
        logger("ERROR", "Could n't save VM's state")

    for task in tasks:   
        task.kill()

    sys.exit(0) 
#    boinc_finish()

###############################################################################
# Standard boilerplate to call the main() function to begin the program.
###############################################################################
# 
if __name__ == '__main__':
    main()

